from langchain_core.prompts.few_shot import FewShotPromptTemplate
from langchain_core.prompts.prompt import PromptTemplate
from open_ai import invoke_open_ai

examples = [
    {
        "question": "谁的寿命更长，张三还是李四？",
        "answer": """
            这里需要跟进问题吗：是的。
            跟进：张三去世时多大？
            中间答案：张三去世时74岁。
            跟进：李四去世时多大？
            中间答案：李四去世时41岁。
            所以最终答案是：张三。
            """,
    },
    {
        "question": "craigslist的创始人是什么时候出生的？",
        "answer": """
            这里需要跟进问题吗：是的。
            跟进：craigslist的创始人是谁？
            中间答案：craigslist由Crag Newmark创立。
            跟进：Crag Newmark是什么时候出生的？
            中间答案：Crag Newmark于1952年12月6日出生。
            所以最终答案是：1952年12月6日
            """,
    },
    {
        "question": "《大白鲨》和《皇家赌场》的导演都来自同一个国家吗？",
        "answer": """
            这里需要跟进问题吗：是的。
            跟进：《大白鲨》的导演是谁？
            中间答案：《大白鲨》的导演是Steven Spielberg。
            跟进：Steven Spielberg来自哪里？
            中间答案：美国。
            跟进：《皇家赌场》的导演是谁？
            中间答案：《皇家赌场》的导演是Martin Campbell。
            跟进：Martin Campbell来自哪里？
            中间答案：新西兰。
            所以最终答案是：不是
            """,
    },
]

example_prompt = PromptTemplate(
    input_variables=["question", "answer"], template="问题：{question}\n{answer}"
)

# 根据下标只取一个实例，转换成参数的形式：
# examples[1]：{"question":"...","answer":"..."}
# **examples[1]：question="...", answer="..."
# print(example_prompt.format(**examples[1]))

prompt = FewShotPromptTemplate(
    examples=examples,
    example_prompt=example_prompt,
    suffix="问题：{input}",
    input_variables=["input"],
)

result = prompt.format(input="张三与李四，谁的寿命更长？分别活了多少岁？")
print(result)
print("\n==============以下是LLM回答================\n")

result1 = invoke_open_ai(prompt, "张三与李四，谁的寿命更长？分别活了多少岁？")
print(result1)
